import torch
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, Sequential
from torch.utils.tensorboard import SummaryWriter

'''
定义了一个基于PyTorch的神经网络模块 Module，继承自 nn.Module。
1. 在 __init__ 方法中，使用 Sequential 定义了卷积层、池化层、展平层和全连接层的序列结构。
2. forward 方法实现了前向传播逻辑，将输入数据通过定义的序列模型进行处理。
3. 创建 Module 实例并打印模型结构，测试输入张量 (64, 3, 32, 32)，输出最终结果的形状。
'''
class Module(nn.Module):
    def __init__(self):
        super().__init__()
        self.module1 = Sequential(
            Conv2d(in_channels=3,out_channels=32,kernel_size=5,stride=1,padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=False),
            Conv2d(in_channels=32,out_channels=32,kernel_size=5,stride=1,padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=False),
            Conv2d(in_channels=32,out_channels=64,kernel_size=5,stride=1,padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=False),
            Flatten(),
            Linear(1024,64),
            Linear(64,10)
        )

    def forward(self,x):
        x = self.module1(x)
        return x

module = Module()
print(module)
input = torch.ones((64,3,32,32))
output = module(input)
print(output.shape)

writer = SummaryWriter("../../logs_seq")
writer.add_graph(module,input)
writer.close()